"""
系统监控数据采集层
使用 psutil 库获取系统信息
"""

import psutil
import time
from typing import Dict, List, Tuple


class SystemMonitor:
    """系统监控数据采集类"""
    
    def __init__(self):
        """初始化系统监控器"""
        self.last_net_io = psutil.net_io_counters()
        self.last_time = time.time()
    
    def get_cpu_info(self) -> Dict[str, float]:
        """获取 CPU 信息"""
        try:
            cpu_percent = psutil.cpu_percent(interval=0.1)
            cpu_count = psutil.cpu_count()
            cpu_freq = psutil.cpu_freq()
            
            return {
                'usage_percent': cpu_percent,
                'core_count': cpu_count,
                'current_freq': cpu_freq.current if cpu_freq else 0,
                'max_freq': cpu_freq.max if cpu_freq else 0
            }
        except Exception as e:
            print(f"获取 CPU 信息失败: {e}")
            return {
                'usage_percent': 0,
                'core_count': 0,
                'current_freq': 0,
                'max_freq': 0
            }
    
    def get_memory_info(self) -> Dict[str, float]:
        """获取内存信息"""
        try:
            memory = psutil.virtual_memory()
            swap = psutil.swap_memory()
            
            return {
                'total': memory.total / (1024**3),  # GB
                'available': memory.available / (1024**3),  # GB
                'used': memory.used / (1024**3),  # GB
                'usage_percent': memory.percent,
                'swap_total': swap.total / (1024**3),  # GB
                'swap_used': swap.used / (1024**3),  # GB
                'swap_percent': swap.percent
            }
        except Exception as e:
            print(f"获取内存信息失败: {e}")
            return {
                'total': 0, 'available': 0, 'used': 0, 'usage_percent': 0,
                'swap_total': 0, 'swap_used': 0, 'swap_percent': 0
            }
    
    def get_disk_info(self) -> Dict[str, float]:
        """获取磁盘信息（主要是 C 盘）"""
        try:
            disk_usage = psutil.disk_usage('C:')
            disk_io = psutil.disk_io_counters()
            
            return {
                'total': disk_usage.total / (1024**3),  # GB
                'used': disk_usage.used / (1024**3),  # GB
                'free': disk_usage.free / (1024**3),  # GB
                'usage_percent': (disk_usage.used / disk_usage.total) * 100,
                'read_bytes': disk_io.read_bytes if disk_io else 0,
                'write_bytes': disk_io.write_bytes if disk_io else 0
            }
        except Exception as e:
            print(f"获取磁盘信息失败: {e}")
            return {
                'total': 0, 'used': 0, 'free': 0, 'usage_percent': 0,
                'read_bytes': 0, 'write_bytes': 0
            }
    
    def get_network_info(self) -> Dict[str, float]:
        """获取网络信息"""
        try:
            current_net_io = psutil.net_io_counters()
            current_time = time.time()
            
            # 计算网络速度（字节/秒）
            time_delta = current_time - self.last_time
            if time_delta > 0:
                bytes_sent_per_sec = (current_net_io.bytes_sent - self.last_net_io.bytes_sent) / time_delta
                bytes_recv_per_sec = (current_net_io.bytes_recv - self.last_net_io.bytes_recv) / time_delta
            else:
                bytes_sent_per_sec = 0
                bytes_recv_per_sec = 0
            
            # 更新上次记录
            self.last_net_io = current_net_io
            self.last_time = current_time
            
            return {
                'bytes_sent': current_net_io.bytes_sent / (1024**2),  # MB
                'bytes_recv': current_net_io.bytes_recv / (1024**2),  # MB
                'packets_sent': current_net_io.packets_sent,
                'packets_recv': current_net_io.packets_recv,
                'send_speed': bytes_sent_per_sec / 1024,  # KB/s
                'recv_speed': bytes_recv_per_sec / 1024   # KB/s
            }
        except Exception as e:
            print(f"获取网络信息失败: {e}")
            return {
                'bytes_sent': 0, 'bytes_recv': 0, 'packets_sent': 0,
                'packets_recv': 0, 'send_speed': 0, 'recv_speed': 0
            }
    
    def get_process_list(self) -> List[Dict]:
        """获取进程列表"""
        try:
            processes = []
            for proc in psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent', 'status']):
                try:
                    proc_info = proc.info
                    # 获取内存使用量（MB）
                    memory_info = proc.memory_info()
                    memory_mb = memory_info.rss / (1024 * 1024)
                    
                    processes.append({
                        'pid': proc_info['pid'],
                        'name': proc_info['name'],
                        'cpu_percent': proc_info['cpu_percent'] or 0,
                        'memory_percent': proc_info['memory_percent'] or 0,
                        'memory_mb': memory_mb,
                        'status': proc_info['status']
                    })
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    # 进程可能已经结束或无权限访问
                    continue
            
            # 按 CPU 使用率排序
            processes.sort(key=lambda x: x['cpu_percent'], reverse=True)
            return processes
            
        except Exception as e:
            print(f"获取进程列表失败: {e}")
            return []
    
    def kill_process(self, pid: int) -> Tuple[bool, str]:
        """结束指定进程"""
        try:
            proc = psutil.Process(pid)
            proc_name = proc.name()
            proc.terminate()
            
            # 等待进程结束
            try:
                proc.wait(timeout=3)
                return True, f"进程 {proc_name} (PID: {pid}) 已成功结束"
            except psutil.TimeoutExpired:
                # 如果进程没有在3秒内结束，强制杀死
                proc.kill()
                return True, f"进程 {proc_name} (PID: {pid}) 已强制结束"
                
        except psutil.NoSuchProcess:
            return False, f"进程 (PID: {pid}) 不存在"
        except psutil.AccessDenied:
            return False, f"没有权限结束进程 (PID: {pid})，请以管理员身份运行"
        except Exception as e:
            return False, f"结束进程失败: {str(e)}"
    
    def get_all_system_info(self) -> Dict:
        """获取所有系统信息"""
        return {
            'cpu': self.get_cpu_info(),
            'memory': self.get_memory_info(),
            'disk': self.get_disk_info(),
            'network': self.get_network_info(),
            'processes': self.get_process_list()
        }


# 测试代码
if __name__ == "__main__":
    monitor = SystemMonitor()
    
    print("=== 系统信息测试 ===")
    info = monitor.get_all_system_info()
    
    print(f"CPU 使用率: {info['cpu']['usage_percent']:.1f}%")
    print(f"内存使用率: {info['memory']['usage_percent']:.1f}%")
    print(f"磁盘使用率: {info['disk']['usage_percent']:.1f}%")
    print(f"网络上传速度: {info['network']['send_speed']:.1f} KB/s")
    print(f"网络下载速度: {info['network']['recv_speed']:.1f} KB/s")
    print(f"进程数量: {len(info['processes'])}")
    
    # 显示前5个占用CPU最高的进程
    print("\n=== CPU 占用最高的前5个进程 ===")
    for i, proc in enumerate(info['processes'][:5]):
        print(f"{i+1}. {proc['name']} (PID: {proc['pid']}) - CPU: {proc['cpu_percent']:.1f}%")